import numpy as np
import numpy.linalg as la
import os
import matplotlib.pyplot as plt
from PIL import Image
import matplotlib.image as mpimg

PATH = os.path.dirname(os.path.abspath(__file__))  # 当前文件夹路径
DATA = os.path.join(PATH, 'data')  # 数据文件夹


def getImgAsMatrix(imgPath):
    return Image.open(imgPath)


def dealImg(imagePath):
    # 将彩色图转为灰度图
    imgM = mpimg.imread(imagePath)
    row, col, channel = imgM.shape
    imgM_gray = np.zeros((row, col))
    for r in range(row):
        for l in range(col):
            imgM_gray[r, l] = 1 / 3 * imgM[r, l, 0] + 1 / 3 * imgM[r, l, 1] + 1 / 3 * imgM[r, l, 2]  # 加权平均方法
    return imgM_gray


def imgSVD(imgM, k=30):
    h, w = imgM.shape[:2]
    U, S, Vt = la.svd(imgM)  # 奇异值分解

    S1 = np.diag(S[:k], 0)  # 奇异值矩阵
    U1 = np.zeros((h, k), float)
    Vt1 = np.zeros((k, w), float)
    U1[:, :] = U[:, :k]
    Vt1[:, :] = Vt[:k, :]
    return U1.dot(S1).dot(Vt1)


if __name__ == '__main__':
    imagePath = os.path.join(DATA, 'whitesnow.jpeg')

    original = getImgAsMatrix(imagePath)

    imgM = dealImg(imagePath)

    newImgM1 = imgSVD(imgM, k=50)

    plt.figure(figsize=(14, 12))

    plt.subplot(221)
    plt.imshow(original)
    plt.title('original')

    plt.subplot(222)
    plt.imshow(Image.fromarray(imgM).convert('RGB'))  # 转为JPEGImageFile进行显示
    plt.title('original gray')

    plt.subplot(223)
    plt.imshow(Image.fromarray(newImgM1).convert('RGB'))
    plt.title('SVD result: k = 50')

    plt.show()
